# -*- coding: utf-8 -*-
# This project is used to ...
#                         ...
# Created by XiangDong Yang on July 7th, 2023
# Copyright (c) 2023 yangxiangdong.cs@aliyun.com

import json
import os

import numpy as np

from environment.umec import Env
from optimize.partition import Welzl, Cluster
from src.optimize.trajectory_planning import TP
from src.optimize.kkt import KKT
from src.optimize.kkt2 import KKT2
from src.optimize.admm import ADMM
from src.optimize.cd import CVX
from src.optimize.other import Ramdom
from src.util.util import DataCollector

if __name__ == '__main__':

    # 0. 获取环境配置
    np.random.seed(0)
    env_config = os.path.join(os.path.dirname(os.getcwd()), 'env.json')

    # 1. 根据配置文件初始化环境，并按照一定聚集性随机初始化 UEs，UAV
    with open(env_config, 'r') as config:
        env = Env(json.load(config))
        env.new_task()

    # 2. 执行 Partition
    # env.partition(Cluster(env=env))
    env.partition(Welzl(radius=env.config.uav_radius))
    # env.plot_distribution()
    env.generate_path()

    # 3. 计算路径
    env.init_meta()

    # 3. 执行 Resource allocation
    collector = DataCollector()

    # kkt_model = KKT2(env)
    # action, bandwidth, frequency = kkt_model.run(collector)
    # collector.write("data_02_01.csv")

    # admm_model = ADMM(env)
    # action, bandwidth, frequency = admm_model.run(DataCollector())

    # cvx_model = CVX(env)
    # action, bandwidth, frequency = cvx_model.run(collector)

    random_model = Ramdom(env)
    action, bandwidth, frequency = random_model.run(collector)

    # for g in env.path:
    #     ues_idx = env.parts_ues[g]
    #     ues_num = int(sum(action[ues_idx]))
    #     collector.collect(g, ues_num)
    #     collector.next()
    # collector.write("data_03_03.csv")
